import os
import hashlib
import json
import ast
import requests
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED
import dotenv
from openai import OpenAI
import time
from pybloom_live import BloomFilter
import gc
import pickle
dotenv.load_dotenv()

class CodeReviewGenerator:
    def __init__(self, api_key,api_url):
        self.api_url = api_url
        self.api_key = api_key
        self.bloom_filter_path = 'code_hashes.bloom'
        self.progress_path = 'generation_progress.json'
        # 尝试从文件加载布隆过滤器，如果不存在则创建新的
        self.code_hashes = self._load_or_create_bloom_filter()
        # 加载已生成的数据量
        self.current_progress = self._load_progress()
        self.error_dist = {
            '逻辑错误': 0.25,
            '异常处理': 0.2,
            '安全漏洞': 0.15,
            '性能反模式': 0.15,
            '代码异味': 0.15,
            '资源泄漏': 0.1
        }

    def _load_or_create_bloom_filter(self):
        """加载或创建布隆过滤器"""
        try:
            if os.path.exists(self.bloom_filter_path):
                print("从本地加载已存在的布隆过滤器...")
                with open(self.bloom_filter_path, 'rb') as f:
                    return pickle.load(f)
            else:
                print("创建新的布隆过滤器...")
                return BloomFilter(capacity=10000, error_rate=0.001)
        except Exception as e:
            print(f"加载布隆过滤器失败: {str(e)}，创建新的过滤器...")
            return BloomFilter(capacity=10000, error_rate=0.001)

    def save_bloom_filter(self):
        """保存布隆过滤器到文件"""
        try:
            with open(self.bloom_filter_path, 'wb') as f:
                pickle.dump(self.code_hashes, f)
            print("布隆过滤器已保存到本地")
        except Exception as e:
            print(f"保存布隆过滤器失败: {str(e)}")

    def _load_progress(self):
        """加载生成进度"""
        try:
            if os.path.exists(self.progress_path):
                with open(self.progress_path, 'r') as f:
                    return json.load(f)
            return {'successful_samples': 0, 'failed_samples': 0}
        except Exception as e:
            print(f"加载进度失败: {str(e)}")
            return {'successful_samples': 0, 'failed_samples': 0}

    def save_progress(self, successful_samples, failed_samples):
        """保存生成进度"""
        try:
            progress = {
                'successful_samples': successful_samples,
                'failed_samples': failed_samples
            }
            with open(self.progress_path, 'w') as f:
                json.dump(progress, f)
        except Exception as e:
            print(f"保存进度失败: {str(e)}")

    def generate_prompt(self, category):
        """动态生成符合分布要求的提示模板"""
        base_prompt = f"""你需要扮演资深Python代码审查员，
                生成包含真实缺陷的代码片段和对应的审查意见，遵循以下规则：
                1. 代码必须包含至少一个{category}类问题
                2. 使用实际业务场景（如：{self._gen_scenario(category)}）
                3. 包含Python特性问题（如：{self._gen_python_spec(category)}）
                4. 输出格式严格遵循JSON，包含instruction,input,output三个键
                example:
                {{
                    "instruction": "Analyze the following code and provide a review:",
                    "input": "code",
                    "output": "review"
                }}
                注意：输出类型instruction,input,output必须是字符串类型，
                input类型为python代码，output类型为代码审查意见。
                除了代码之外的内容，要使用中文输出。
                """

        return base_prompt

    def _gen_scenario(self, category):
        # 场景动态生成逻辑
        scenarios = {
            '安全漏洞': ['API端点输入验证', 'JWT令牌处理', '文件上传检测'],
            '资源泄漏': ['数据库连接池管理', '文件流操作', '缓存系统交互'],
            '逻辑错误': ['订单处理流程', '支付状态同步', '库存管理系统'],
            '异常处理': ['第三方API调用', '分布式任务处理', '消息队列消费'],
            '性能反模式': ['大数据批处理', '实时数据分析', '缓存更新策略'],
            '代码异味': ['用户认证模块', '日志记录系统', '配置管理服务']
        }
        return scenarios.get(category, ['数据处理管道', 'Web请求处理', '异步任务处理'])[0]

    def _gen_python_spec(self, category):
        # Python特性选择逻辑
        specs = {
            '代码异味': ['不当的魔法方法使用', 'PEP8违规', '过度复杂的条件语句', '重复代码块'],
            '性能反模式': ['低效的迭代操作', '不必要的对象深拷贝', '频繁的字符串拼接', '大量小对象创建'],
            '逻辑错误': ['边界条件处理不当', '不正确的异常捕获顺序', '死锁风险'],
            '异常处理': ['裸露的异常捕获', '异常屏蔽', '自定义异常使用不当'],
            '安全漏洞': ['SQL注入风险', 'XSS攻击', '不安全的反序列化'],
            '资源泄漏': ['未关闭的文件句柄', '数据库连接未释放', '线程资源未清理']
        }
        return specs.get(category, '上下文管理器误用')

    def call_api(self, prompt):
        """调用API"""
        try:
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    url = "http://localhost:11434/api/generate"
                    data = {
                        "model": "qwen2.5-coder:1.5b",
                        "prompt": prompt,
                        "stream": False,
                        "format": "json",
                        "temperature": 0.6
                    }
                    # 添加超时设置
                    response = requests.post(url, json=data, timeout=30)
                    if response.status_code != 200:
                        print(f"API响应状态码: {response.status_code}")
                        raise Exception(f"API调用失败: {response.text}")
                        
                    data = json.loads(response.text)
                    result = data.get("response")
                    if not result:
                        raise Exception("API返回空响应")
                    return result
                    
                except requests.Timeout:
                    print(f"API调用超时 (尝试 {attempt + 1}/{max_retries})")
                except requests.RequestException as e:
                    print(f"请求异常: {str(e)} (尝试 {attempt + 1}/{max_retries})")
                except Exception as e:
                    print(f"其他错误: {str(e)} (尝试 {attempt + 1}/{max_retries})")
                
                if attempt < max_retries - 1:
                    sleep_time = 2 ** attempt  # 指数退避：1, 2, 4秒
                    print(f"等待 {sleep_time} 秒后重试...")
                    time.sleep(sleep_time)
            
        except Exception as e:
            print(f"API调用最终失败: {str(e)}")
            print(f"提示文本: {prompt}")
        return None

    def validate_code(self, code_snippet):
        """执行代码验证"""
        try:
            ast.parse(code_snippet)
            return True
        except SyntaxError:
            return False

    def check_uniqueness(self, code):
        """使用布隆过滤器进行重复检测"""
        code_hash = hashlib.md5(code.encode()).hexdigest()
        if code_hash in self.code_hashes:
            return False
        self.code_hashes.add(code_hash)
        return True

    def generate_dataset(self, num_samples=1000, batch_size=100):
        try:
            output_file = 'python_code_review_dataset.jsonl'
            successful_samples = self.current_progress['successful_samples']
            failed_samples = self.current_progress['failed_samples']
            remaining_samples = num_samples - successful_samples
            
            if remaining_samples <= 0:
                print(f"已完成目标数量（{num_samples}条），无需继续生成")
                return successful_samples
            
            print(f"从第 {successful_samples + 1} 条开始继续生成，目标总数：{num_samples}")
            
            with tqdm(total=remaining_samples, desc="生成数据集") as pbar:
                for batch_start in range(0, remaining_samples, batch_size):
                    batch_end = min(batch_start + batch_size, remaining_samples)
                    remaining_in_batch = batch_end - batch_start
                    
                    with ThreadPoolExecutor(max_workers=5) as executor:
                        futures = set()
                        
                        with open(output_file, 'a', encoding='utf-8') as f:
                            while remaining_in_batch > 0 or futures:
                                # 提交新任务
                                while remaining_in_batch > 0 and len(futures) < 5:
                                    category = self._select_category(list(self.error_dist.items()))
                                    prompt = self.generate_prompt(category)
                                    future = executor.submit(self.process_sample, prompt)
                                    futures.add(future)
                                    remaining_in_batch -= 1
                                
                                # 等待并处理完成的任务
                                if futures:
                                    done, pending = wait(futures, timeout=1, return_when=FIRST_COMPLETED)
                                    # 更新 futures 集合
                                    futures = pending
                                    
                                    for future in done:
                                        try:
                                            result = future.result()
                                            if result:
                                                f.write(json.dumps(result, ensure_ascii=False) + '\n')
                                                f.flush()
                                                successful_samples += 1
                                            else:
                                                failed_samples += 1
                                        except Exception as e:
                                            print(f"\n任务执行失败: {str(e)}")
                                            failed_samples += 1
                                        finally:
                                            pbar.update(1)
                                            pbar.set_postfix({
                                                '成功': successful_samples,
                                                '失败': failed_samples,
                                                '成功率': f"{(successful_samples/(successful_samples+failed_samples)*100):.1f}%"
                                            })
                    
                    # 每个批次完成后进行垃圾回收
                    gc.collect()
                    
                    # 每个批次结束后保存布隆过滤器
                    self.save_bloom_filter()
                    
                    # 打印批次统计
                    print(f"\n批次完成:")
                    print(f"- 已生成样本数: {successful_samples}")
                    print(f"- 失败样本数: {failed_samples}")
                    print(f"- 总成功率: {(successful_samples/(successful_samples+failed_samples)*100):.1f}%")
                    print(f"- 距离目标还需: {num_samples - successful_samples} 条")
            
            return successful_samples
            
        except KeyboardInterrupt:
            print("\n检测到程序中断，正在保存进度...")
            self.save_progress(successful_samples, failed_samples)
            self.save_bloom_filter()
            raise
        except Exception as e:
            print(f"\n发生错误: {str(e)}，正在保存进度...")
            self.save_progress(successful_samples, failed_samples)
            self.save_bloom_filter()
            raise
        finally:
            # 确保在程序结束时保存布隆过滤器
            self.save_bloom_filter()

    def process_sample(self, prompt):
        """处理单个样本"""
        retry = 3
        while retry > 0:
            try:
                response = self.call_api(prompt)
                if response:
                    try:
                        data = json.loads(response)
                        if self._validate_sample(data):
                            return data
                    except json.JSONDecodeError as e:
                        print(f"JSON解析错误: {str(e)}, 剩余重试次数: {retry-1}")
                retry -= 1
            except Exception as e:
                print(f"处理样本异常: {str(e)}, 剩余重试次数: {retry-1}")
                retry -= 1
        return None

    def _validate_sample(self, data):
        """样本验证"""
        required_keys = {'instruction', 'input', 'output'}
        if not all(k in data for k in required_keys):
            return False

        code = data['input'].split('```')[1][len('python'):] if '```' in data['input'] else data['input']
        
        return (
            self.validate_code(code) and
            self.check_uniqueness(code) 
        )

    def _check_security_rules(self, output):
        """安全检查
        Args:
            output: 代码审查输出文本
            category: 错误类型，如果指定则使用对应的关键词集合
        Returns:
            bool: 是否通过安全检查
        """
        # 基础关键词集合
        base_keywords = {'CWE-', 'PEP8', '建议使用', '应改为', '需要修改', '存在风险'}
        
        # 针对不同错误类型的特定关键词
        category_keywords = {
            '逻辑错误': {'边界条件', '逻辑判断', '控制流程', '业务规则'},
            '异常处理': {'try-except', '异常捕获', '错误处理', '防御性编程'},
            '安全漏洞': {'注入攻击', 'SQL注入', 'XSS', '越权访问', 'CSRF'},
            '性能反模式': {'时间复杂度', '空间复杂度', '内存泄漏', '性能优化'},
            '代码异味': {'代码重复', '命名规范', '函数过长', '耦合度'},
            '资源泄漏': {'文件句柄', '数据库连接', '内存资源', '线程池'}
        }
        
        # 合并关键词集合
        check_keywords = base_keywords.copy()  # 创建副本避免修改原始集合
        
        # 将所有category_keywords中的值（集合）合并到check_keywords中
        for keyword_set in category_keywords.values():
            check_keywords.update(keyword_set)
        
        return any(kw in output for kw in check_keywords)

    def _select_category(self, dist):
        # 基于概率分布选择错误类型
        import random
        r = random.random()
        cumulative = 0.0
        for category, prob in dist:
            cumulative += prob
            if r < cumulative:
                return category
        return dist[-1][0]

# 使用示例
if __name__ == "__main__":
    api_key = os.getenv("DEEPSEEK_API_KEY")
    api_url = os.getenv("DEEPSEEK_API_URL")
    generator = CodeReviewGenerator(api_key, api_url)
    # 生成5000条数据，每批次处理100条
    total_samples = generator.generate_dataset(5000, batch_size=100)
    print(f"总共生成 {total_samples} 条有效数据")